package client

import (
	"context"
	"sync"
	"time"

	"github.com/coreos/etcd/mvcc/mvccpb"

	"github.com/coreos/etcd/clientv3"
)

/**
服务发现步骤：
1、connect etcd
2、get value
3、watch key
*/

const ETCD_TIMEOUT = 3 * time.Second

type EtcdDiscover struct {
	cli        *clientv3.Client
	endpoints  []string
	ServerList sync.Map
}

// NewEtcdDiscover new一个服务发现者
func NewEtcdDiscover(endpoints []string) (etcd *EtcdDiscover, err error) {
	config := clientv3.Config{
		Endpoints:   endpoints,
		DialTimeout: ETCD_TIMEOUT,
	}
	cli, err := clientv3.New(config)
	if err != nil {
		return nil, err
	}
	ts := &EtcdDiscover{
		cli:        cli,
		endpoints:  config.Endpoints,
		ServerList: sync.Map{},
	}
	return ts, nil
}

// GetValueByPrefix 获取值根据前缀
func (e *EtcdDiscover) GetValueByPrefix(prefix string) error {
	ctx, cancel := context.WithTimeout(context.Background(), ETCD_TIMEOUT)
	res, err := e.cli.Get(ctx, prefix, clientv3.WithPrefix())
	cancel()
	if err != nil {
		return err
	}
	if res.Count > 0 {
		for _, v := range res.Kvs {
			e.ServerList.Store(string(v.Key), string(v.Value))
		}
	}
	return nil
}

// WatchPrefix 监听
func (e *EtcdDiscover) WatchPrefix(prefix string) {
	watchChan := e.cli.Watch(context.TODO(), prefix, clientv3.WithPrefix())
	for rch := range watchChan {
		for _, event := range rch.Events {
			switch event.Type {
			case mvccpb.PUT:
				e.Add(string(event.Kv.Key), string(event.Kv.Value))
			case mvccpb.DELETE:
				e.Del(string(event.Kv.Key))
			}
		}
	}
}

func (e *EtcdDiscover) Add(k, v string) {
	e.ServerList.Store(k, v)
}

func (e *EtcdDiscover) Del(k string) {
	e.ServerList.Delete(k)
}
